Sys.setenv(OPENAI_API_KEY = # INSERT KEY HERE
)

transcripts_with_embeddings <- read_csv("data/cleaned/transcripts_with_embeddings.csv") %>%
  mutate(speech_english = quote) %>%
  filter(!(who_speaking %in% c("Lead", "Other")))

transcripts_for_morality_coding <- transcripts_with_embeddings %>%
  select(group_id, pair_id, who_speaking_tr, speech_english) %>%
  mutate(speech_english = paste0(who_speaking_tr, ": ", speech_english)) %>%
  group_by(group_id, pair_id) %>%
  summarise(
    speech_english = paste(speech_english, collapse = "\n")
  ) %>%
  ungroup()

# Constants - Changed file extensions from .rds to .csv
SAVE_PATH_BASE <- "data/morality_analysis/"
SAVE_PATH_EXPANDED <- paste0(SAVE_PATH_BASE, "expanded_dataset.csv")
SAVE_PATH_COMPLETED <- paste0(SAVE_PATH_BASE, "completed_rows.csv")
SAVE_PATH_FINAL <- paste0(SAVE_PATH_BASE, "final_analysis.csv")

remove_existing <- FALSE

if (remove_existing) {
  #   Remove existing files
  file.remove(SAVE_PATH_EXPANDED, SAVE_PATH_COMPLETED, SAVE_PATH_FINAL)
}

# Create safe versions that always return exactly one value
safe_get_rating <- function(x) {
  result <- tryCatch(
  {
    if (is.null(x$rating) || length(x$rating) == 0) NA_real_
    else x$rating
  },
    error = function(e) NA_real_
  )
  if (length(result) == 0) NA_real_
  else result[1]  # Take first value if multiple are returned
}

safe_get_explanation <- function(x) {
  result <- tryCatch(
  {
    if (is.null(x$explanation) || length(x$explanation) == 0) NA_character_
    else x$explanation
  },
    error = function(e) NA_character_
  )
  if (length(result) == 0) NA_character_
  else result[1]  # Take first value if multiple are returned
}

sd_na <- function (x) {
  l <- length(x[!is.na(x)])
  if (l > 0)
    sd(x, na.rm = T)
  else if (l == 0)
    NA
  else stop("something's wrong")
}


# Create directory if it doesn't exist
dir.create(SAVE_PATH_BASE, recursive = TRUE, showWarnings = FALSE)

# Function to generate random moral topics (unchanged)
generate_random_moral_topics <- function() {
  all_topics <- c(
    "Fairness",
    "Rights",
    "Justice",
    "Giving people opportunities",
    "Ethical decision-making",
    "What is 'right' or 'wrong'",
    "Responsibility and accountability",
    "Personal autonomy",
    "Care and compassion",
    "Human dignity",
    "Truth and honesty"
  )
  sample(all_topics, 3)
}

# Function to generate random prompt (unchanged)
generate_random_prompt <- function() {
  include_reflect_both <- sample(c(TRUE, FALSE), 1)
  include_explicit_refs <- sample(c(TRUE, FALSE), 1)
  moral_topics <- generate_random_moral_topics()

  base_prompt <- "
BACKGROUND:
    Below is a transcript from a group discussion in Chennai, India, in which 3 participants (R1, R2, R3) discussed who they would prefer to hire to deliver groceries to their home. Participants had to choose between option A or B.
    They were shown photos of the two delivery workers, one of whom was male, and one of whom was either male, female, or transgender.
    Each discussion discussed multiple sets of options, changing the choice of Option A and Option B each time.
    The grocery items on offer were: Aachi masala spice, tea powder, and ghee. Option A and B may have offered the same set of items, or different sets of items.
    They were also in some cases given information about:
     - the languages spoken by the delivery workers (Tamil only, Tamil and English)
     - the delivery workers' experience
     - how many deliveries they completed in a training task
     Participants were asked to discuss which option they prefer, and why, and then make a collective choice between the two options.

TASK: Analyze the transcript by rating how much the discussants invoke morality in their decision-making process on a scale from 0 to 1.
Rating Scale Details:
0: No moral considerations mentioned
1: Morality is the central theme and dominant focus\n"

  reflect_both_text <- if(include_reflect_both) {
    "The rating should reflect both:
- The frequency of mentions of moral considerations
- The intensity/emphasis of moral considerations when they appear\n"
  } else ""

  explicit_refs_text <- if(include_explicit_refs) {
    "Explicit moral references (e.g., direct mentions of fairness, rights, justice, or 'right thing to do') should be weighted more heavily than implicit ones.\n"
  } else ""

  moral_topics_text <- sprintf("
Moral considerations include topics like:
- %s
- %s
- %s

Note: Simply mentioning demographic characteristics (like gender) without connecting them to moral concepts does not count as invoking morality.

Return your analysis in this JSON format:
{
    \"explanation\": \"{EXPLANATION OF YOUR RATING}\",
    \"rating\": {YOUR RATING}
}", moral_topics[1], moral_topics[2], moral_topics[3])

  full_prompt <- paste0(
    base_prompt,
    reflect_both_text,
    explicit_refs_text,
    moral_topics_text
  )

  list(
    prompt = full_prompt,
    config = list(
      include_reflect_both = include_reflect_both,
      include_explicit_refs = include_explicit_refs,
      moral_topics = moral_topics
    )
  )
}

# Modified llm_analyse function (unchanged)
llm_analyse <- function(text, prompt, temperature = runif(1, 0.1, 0.9), .i = NULL, model = "gpt-4o-mini") {
  if (!is.null(.i)) {
    print(str_glue("\nProcessing Transcript: {.i}"))
  }

  out <- create_chat_completion(
    model = model,
    messages = list(
      list(
        "role" = "system",
        "content" = prompt
      ),
      list(
        "role" = "user",
        "content" = str_glue("TRANSCRIPT: {text}")
      )
    ),
    temperature = temperature
  )

  out_clean <- out$choices$message.content[1] %>%
    str_replace_all("```(json)?", "") %>%
    fromJSON()

  out_clean$temperature <- temperature
  return(out_clean)
}


# Modified function to load completed analyses
load_completed_analyses <- function() {
  if (file.exists(SAVE_PATH_COMPLETED)) {
    read_csv(SAVE_PATH_COMPLETED) %>%
      mutate(
        analysis = map(analysis, ~fromJSON(.x)),
        moral_topics = map(moral_topics, ~fromJSON(.x)),
        prompt_info = map(prompt_info, ~fromJSON(.x))
      ) %>%
      print
  } else {
    tibble()
  }
}

# Modified function to save completed analyses
save_completed_analyses <- function(data) {
  data %>%
    mutate(
      analysis = map_chr(analysis, ~toJSON(.x, auto_unbox = TRUE)),
      moral_topics = map_chr(moral_topics, ~toJSON(.x, auto_unbox = TRUE)),
      prompt_info = map_chr(prompt_info, ~toJSON(.x, auto_unbox = TRUE))
    ) %>%
    write_csv(SAVE_PATH_COMPLETED)
}

# Modified main analysis function
run_morality_analysis <- function(transcripts_df, n_reps_for_rows = 50, model = "gpt-4o-mini") {
  # Step 1: Create expanded dataset with repetitions and random prompts


  if (file.exists(SAVE_PATH_EXPANDED)) {
    message("Using existing expanded dataset...")
    expanded_data <- read_csv(SAVE_PATH_EXPANDED) %>%
      mutate(
        moral_topics = map(moral_topics, ~fromJSON(.x)),
        prompt_info = map(prompt_info, ~fromJSON(.x))
      )
  } else {
    message("Creating expanded dataset...")
    expanded_data <- transcripts_df %>%
      crossing(rep_number = 1:n_reps_for_rows) %>%
      rowwise() %>%
      mutate(
        prompt_info = list(generate_random_prompt()),
        prompt = prompt_info$prompt,
        include_reflect_both = prompt_info$config$include_reflect_both,
        include_explicit_refs = prompt_info$config$include_explicit_refs,
        moral_topics = list(prompt_info$config$moral_topics),
        temperature = runif(1, 0.1, 0.9)
      ) %>%
      ungroup()

    # Save expanded dataset (converting list columns to JSON)
    expanded_data %>%
      mutate(
        moral_topics = map_chr(moral_topics, ~toJSON(.x, auto_unbox = TRUE)),
        prompt_info = map_chr(prompt_info, ~toJSON(.x, auto_unbox = TRUE))
      ) %>%
      write_csv(SAVE_PATH_EXPANDED)

  }

  # Step 2: Load completed analyses
  completed_analyses <- load_completed_analyses()

  # Step 3: Process remaining rows
  message("Processing rows...")
  for (i in seq_len(nrow(transcripts_df))) {
    current_group_id <- transcripts_df$group_id[i]
    current_pair_id <- transcripts_df$pair_id[i]

    # Check if this row is already completed
    if (nrow(completed_analyses) > 0 &&
      any(completed_analyses$group_id == current_group_id &
            completed_analyses$pair_id == current_pair_id)) {
      message(sprintf("Skipping completed row %d (group_id: %s, pair_id: %s)",
                      i, current_group_id, current_pair_id))
      next
    }

    # Get current row's data
    current_rows <- expanded_data %>%
      filter(group_id == current_group_id, pair_id == current_pair_id)

    message(sprintf("Processing row %d of %d (group_id: %s, pair_id: %s)",
                    i, nrow(transcripts_df), current_group_id, current_pair_id))

    # Process all repetitions for current row
    results <- current_rows %>%
      rowwise() %>%
      mutate(
        analysis = list(safely(llm_analyse)(
          text = speech_english,
          prompt = prompt,
          temperature = temperature,
          model = model,
          .i = sprintf("%s-%s-%d", group_id, pair_id, rep_number)
        )$result)
      ) %>%
      ungroup()

    # Add to completed analyses
    completed_analyses <- bind_rows(completed_analyses, results)

    # Save after each original row is completed
    save_completed_analyses(completed_analyses)

    message(sprintf("Completed row %d", i))
  }

  # completed_analyses %>%


  # Step 4: Calculate summary statistics
  message("Calculating summary statistics...")
  final_analysis <- completed_analyses %>%
    mutate(
      morality_rating = map_dbl(analysis, safe_get_rating),
      explanation = map_chr(analysis, safe_get_explanation)
    ) %>%
    group_by(group_id, pair_id) %>%
    summarise(
      mean_morality_rating = trackr::mean_na(morality_rating),
      sd_morality_rating = sd_na(morality_rating),
      n_observations = sum(!is.na(morality_rating)),
      .groups = "drop"
    )

  # Save final analysis
  write_csv(final_analysis, SAVE_PATH_FINAL)

  # Return all datasets
  list(
    expanded_data = expanded_data,
    completed_analyses = completed_analyses,
    final_analysis = final_analysis
  )
}

# Run the analysis
set.seed(123)
results <- run_morality_analysis(transcripts_for_morality_coding,
                                 n_reps_for_rows = 50,
                                 model = "gpt-4o-mini")


results$final_analysis %>%
  hist_basic(mean_morality_rating)

results$completed_analyses %>%
  mutate(
    explanation = map_chr(analysis, "explanation"),
    rating = map_dbl(analysis, "rating")
  ) %>%
  select(group_id, pair_id, speech_english, rating, explanation) %>%
  hist_basic(rating)


final_analysis %>%
  hist_basic(mean_morality_rating)


final_analysis <- completed_analyses %>%
  mutate(
    morality_rating = map_dbl(analysis, "rating"),
    explanation = map_chr(analysis, "explanation")
  ) %>%
  group_by(group_id, pair_id) %>%
  summarise(
    mean_morality_rating = mean(morality_rating),
    sd_morality_rating = sd(morality_rating),
    n_observations = n(),
    .groups = "drop"
  )